#!/usr/bin/env python
# -*- coding: utf-8 -*-

from typing import List, Dict, Any
import asyncio

from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.future import select

from app.core.logging import get_logger
from app.core.utils import validate_collection_name
from app.core.exceptions import VectorDBError
from app.models.database import Project
from app.repositories.mysql import get_db_session
from app.repositories.milvus import milvus_client
from app.repositories.vector import VectorRepository

# 创建日志记录器
logger = get_logger("vector_init")


async def initialize_vector_collections() -> None:
    """初始化所有项目的向量集合
    
    检查所有项目，确保每个项目都有必要的向量集合
    如果集合不存在，则创建
    """
    try:
        # 连接到Milvus
        await milvus_client.connect()
        
        # 创建向量仓库
        vector_repo = VectorRepository(milvus_client)
        
        # 获取所有项目
        projects = []
        async for session in get_db_session():
            result = await session.execute(select(Project))
            projects = result.scalars().all()
            break  # 只需要执行一次查询
        
        if not projects:
            logger.warning("没有找到任何项目，无需初始化向量集合")
            return
        
        # 为每个项目初始化向量集合
        for project in projects:
            project_id = project.id
            
            logger.info(f"正在为项目 {project_id} 初始化向量集合")
            try:
                # 使用向量仓库的方法确保集合存在
                await vector_repo.ensure_collections(project_id)
            except Exception as e:
                logger.error(f"为项目 {project_id} 初始化向量集合失败: {str(e)}")
                # 继续处理其他项目，不中断流程
        
        logger.info("所有项目的向量集合初始化完成")
    except Exception as e:
        logger.error(f"初始化向量集合失败: {str(e)}")
        raise